#include <memory>
#include <rclcpp/rclcpp.hpp>
#include <moveit/move_group_interface/move_group_interface.h>
#include <moveit/planning_scene_interface/planning_scene_interface.h>
#include "tf2/LinearMath/Quaternion.h"
#include <thread>
#include <chrono>
#include <tf2_ros/buffer.h>
#include <tf2_ros/transform_listener.h>

#define M_PI 3.14159265358979323846

using namespace std::chrono_literals;



// 笛卡尔路径规划辅助函数
bool move_l(
  moveit::planning_interface::MoveGroupInterface& move_group,
  geometry_msgs::msg::Pose target_pose,
  double step_size = 0.01,
  double jump_threshold = 0.0)
{
  std::vector<geometry_msgs::msg::Pose> waypoints;
  waypoints.push_back(target_pose);

  moveit_msgs::msg::RobotTrajectory trajectory;
  const double fraction = move_group.computeCartesianPath(
      waypoints,   // 路径点集合
      step_size,   // 笛卡尔路径步长（米）
      jump_threshold,  // 跳跃阈值（设为0禁用）
      trajectory
  );

  if (fraction >= 0.9) { // 允许部分成功
      moveit::planning_interface::MoveGroupInterface::Plan plan;
      plan.trajectory_ = trajectory;
      move_group.execute(plan);
      return true;
  } else {
      RCLCPP_ERROR(rclcpp::get_logger("rclcpp"), 
          "笛卡尔路径规划失败 (完成度: %.2f%%)", fraction * 100.0);
      return false;
  }
}

void zhua_pose(  moveit::planning_interface::MoveGroupInterface& arm,
  moveit::planning_interface::MoveGroupInterface& gripper,
  geometry_msgs::msg::Pose target_pose)
{

  // 1. 回到home点
  RCLCPP_INFO_STREAM(rclcpp::get_logger("rclcpp"),"Moving to pose: Home");
  arm.setNamedTarget("home");
  arm.move();

  // 2. 移动到物体上方
  RCLCPP_INFO(rclcpp::get_logger("rclcpp"), "Moving above the object");
  target_pose.position.z += 0.42;
  if (!move_l(arm, target_pose)) {
    RCLCPP_ERROR(rclcpp::get_logger("rclcpp"), "直线运动失败，终止程序");
  }

  std::this_thread::sleep_for(1s);

  // 设置一个比例因子以选择性地降低最大关节速度限制，可取值为(0,1]
  arm.setMaxVelocityScalingFactor(0.1);

  // 3. 移动到物体处
  RCLCPP_INFO(rclcpp::get_logger("rclcpp"), "Moving to the object");
  target_pose.position.z -= 0.22;
  if (!move_l(arm, target_pose)) {
    RCLCPP_ERROR(rclcpp::get_logger("rclcpp"), "直线下降运动失败，终止程序");
  }

  std::this_thread::sleep_for(1s);

  // 4. 闭合夹爪
  RCLCPP_INFO(rclcpp::get_logger("rclcpp"), "Closing gripper");
  gripper.setNamedTarget("close");
  gripper.move();

  std::this_thread::sleep_for(1s);

  // 5. 回到物体上方
  RCLCPP_INFO(rclcpp::get_logger("rclcpp"), "Moving above the object");
  target_pose.position.z += 0.22;
  if (!move_l(arm, target_pose)) {
    RCLCPP_ERROR(rclcpp::get_logger("rclcpp"), "直线上升运动失败，终止程序");
  }
  std::this_thread::sleep_for(1s);

  // 6. 回到home点
  RCLCPP_INFO_STREAM(rclcpp::get_logger("rclcpp"),"Moving to pose: Home");
  arm.setNamedTarget("home");
  arm.move();
}

void fang_pose(  moveit::planning_interface::MoveGroupInterface& arm,
  moveit::planning_interface::MoveGroupInterface& gripper,
  geometry_msgs::msg::Pose target_pose)
{

  // 1. 回到home点
  RCLCPP_INFO_STREAM(rclcpp::get_logger("rclcpp"),"Moving to pose: Home");
  arm.setNamedTarget("home");
  arm.move();

  // 2. 移动到物体上方
  RCLCPP_INFO(rclcpp::get_logger("rclcpp"), "Moving above the object");
  target_pose.position.z += 0.42;
  if (!move_l(arm, target_pose)) {
    RCLCPP_ERROR(rclcpp::get_logger("rclcpp"), "直线运动失败，终止程序");
  }

  std::this_thread::sleep_for(1s);

  // 设置一个比例因子以选择性地降低最大关节速度限制，可取值为(0,1]
  arm.setMaxVelocityScalingFactor(0.1);

  // 3. 移动到物体处
  RCLCPP_INFO(rclcpp::get_logger("rclcpp"), "Moving to the object");
  target_pose.position.z -= 0.22;
  if (!move_l(arm, target_pose)) {
    RCLCPP_ERROR(rclcpp::get_logger("rclcpp"), "直线下降运动失败，终止程序");
  }

  std::this_thread::sleep_for(1s);

  // 4. 开放夹爪
  RCLCPP_INFO(rclcpp::get_logger("rclcpp"), "Opening gripper");
  gripper.setNamedTarget("open");
  gripper.move();

  std::this_thread::sleep_for(1s);

  // 5. 回到物体上方
  RCLCPP_INFO(rclcpp::get_logger("rclcpp"), "Moving above the object");
  target_pose.position.z += 0.22;
  if (!move_l(arm, target_pose)) {
    RCLCPP_ERROR(rclcpp::get_logger("rclcpp"), "直线上升运动失败，终止程序");
  }
  std::this_thread::sleep_for(1s);

  // 6. 回到home点
  RCLCPP_INFO_STREAM(rclcpp::get_logger("rclcpp"),"Moving to pose: Home");
  arm.setNamedTarget("home");
  arm.move();
}

int main(int argc, char * argv[])
{
  // Initialize ROS and create the Node
  rclcpp::init(argc, argv);
  auto const node = std::make_shared<rclcpp::Node>(
    "aruco_pick_place_demo",
    rclcpp::NodeOptions().automatically_declare_parameters_from_overrides(true)
  );

  // Create a ROS logger
  auto const logger = rclcpp::get_logger("aruco_pick_place_demo");

  // We spin up a SingleThreadedExecutor so MoveItVisualTools interact with ROS
  rclcpp::executors::SingleThreadedExecutor executor;
  executor.add_node(node);
  auto spinner = std::thread([&executor]() { executor.spin(); });

  // Create the MoveIt MoveGroup Interface
  using moveit::planning_interface::MoveGroupInterface;
  auto arm = MoveGroupInterface(node, "ur_manipulator");
  auto gripper = MoveGroupInterface(node, "gripper");
  auto tf_buffer = tf2_ros::Buffer(node->get_clock());
  auto tf_listener = tf2_ros::TransformListener(tf_buffer, node);

  // 获取xarm规划组的规划参考坐标系
  std::string planning_frame = arm.getPlanningFrame();
  RCLCPP_INFO_STREAM(rclcpp::get_logger("rclcpp"), "Planning frame : "<< planning_frame);
  // 获取末端执行器的link
  std::string eef_link = arm.getEndEffectorLink();
  RCLCPP_INFO_STREAM(rclcpp::get_logger("rclcpp"), "End effector link : "<< eef_link);
  // 若allow_replanning()参数为True，则MoveIt!在一次规划失败后进行重新规划
  arm.allowReplanning(true);
  // 设置运动到目标时的位置(单位：米)和姿态的容忍误差(单位：弧度)
  arm.setGoalPositionTolerance(0.02);
  arm.setGoalOrientationTolerance(0.03);
  // 设置一个比例因子以选择性地降低最大关节速度限制，可取值为(0,1]
  arm.setMaxVelocityScalingFactor(0.5);

  while (rclcpp::ok()) 
  {
  geometry_msgs::msg::TransformStamped transform;
  bool transform_received = false;

  while (!transform_received)
  {
    try
    {
        rclcpp::Time now = node->get_clock()->now();
        transform = tf_buffer.lookupTransform("base_link","aruco_marker_frame", tf2::TimePointZero);
        transform_received = true; // 如果lookupTransform成功，设置为true以退出循环
    }
    catch(tf2::TransformException &ex)
    {
        RCLCPP_ERROR(rclcpp::get_logger("rclcpp"), "没有找到aruco码, TF lookup failed: %s", ex.what());
        // 可以在这里添加一些等待时间，避免频繁请求
        std::this_thread::sleep_for(std::chrono::seconds(1));
    }
  }


  RCLCPP_INFO_STREAM(logger, "Transform: " << transform.transform.translation.x << ", "
                                             << transform.transform.translation.y << ", "
                                             << transform.transform.translation.z);



  RCLCPP_INFO_STREAM(rclcpp::get_logger("rclcpp"),"Moving to pose: Home");
  arm.setNamedTarget("home");
  arm.move();
  gripper.setNamedTarget("open");
  gripper.move();
  std::this_thread::sleep_for(1s);
  
  // 获取末端执行器当前的位姿
  geometry_msgs::msg::Pose current_pose = arm.getCurrentPose().pose;
  // 获取机械臂六个关节当前的位置
  RCLCPP_INFO(rclcpp::get_logger("rclcpp"), "Robot pose x : %.2fm; y : %.2fm; z : %.2fm", 
            current_pose.position.x, current_pose.position.y, current_pose.position.z);
    
    // 提取marker的姿态和位置
    auto marker_pose = transform.transform;
    tf2::Quaternion marker_quat(
        marker_pose.rotation.x,
        marker_pose.rotation.y,
        marker_pose.rotation.z,
        marker_pose.rotation.w);

    // 计算末端执行器姿态 (保持当前姿态，仅更新z轴旋转)
    geometry_msgs::msg::Pose target_pose;
    target_pose.position.x = marker_pose.translation.x;
    target_pose.position.y = marker_pose.translation.y;
    target_pose.position.z = marker_pose.translation.z;

    // 获取marker的z轴旋转角度
    double roll, pitch, yaw;
    tf2::Matrix3x3(marker_quat).getRPY(roll, pitch, yaw);

    // 获取当前姿态的roll和pitch
    tf2::Quaternion current_quat(
        current_pose.orientation.x,
        current_pose.orientation.y,
        current_pose.orientation.z,
        current_pose.orientation.w);

    double current_roll, current_pitch, current_yaw;
    tf2::Matrix3x3(current_quat).getRPY(current_roll, current_pitch, current_yaw);

    // 保持当前roll和pitch，仅使用marker的yaw
    tf2::Quaternion end_effector_quat;
    end_effector_quat.setRPY(current_roll, current_pitch, yaw);

    target_pose.orientation.x = end_effector_quat.x();
    target_pose.orientation.y = end_effector_quat.y();
    target_pose.orientation.z = end_effector_quat.z();
    target_pose.orientation.w = end_effector_quat.w();

    RCLCPP_INFO(rclcpp::get_logger("rclcpp"), "Target pose x : %.2fm; y : %.2fm; z : %.2fm; RZ : %.2fm 度", 
    target_pose.position.x, target_pose.position.y, target_pose.position.z, yaw*180/M_PI);

    // 1.抓取物体
    zhua_pose(arm, gripper, target_pose);
    std::this_thread::sleep_for(1s);

    //2.计算放置点位置
    current_pose.position.y = 0.40;
    current_pose.position.z = target_pose.position.z;
    std::this_thread::sleep_for(1s);

    // 3.放置物体
    fang_pose(arm, gripper, current_pose);
    std::this_thread::sleep_for(1s);

    //4.抓取物体
    zhua_pose(arm, gripper, current_pose);
    std::this_thread::sleep_for(1s);

    //5.放置物体
    fang_pose(arm, gripper, target_pose);
    std::this_thread::sleep_for(1s);
    
    std::this_thread::sleep_for(15s);

  }
    // 设置一个比例因子以选择性地降低最大关节速度限制，可取值为(0,1]
  arm.setMaxVelocityScalingFactor(0.5);
  
  RCLCPP_INFO_STREAM(rclcpp::get_logger("rclcpp"),"Moving to pose: Home");
  arm.setNamedTarget("home");
  arm.move();

  std::this_thread::sleep_for(1s);

  // 获取末端执行器当前的位姿
  geometry_msgs::msg::Pose test_pose = arm.getCurrentPose().pose;
  // 获取机械臂六个关节当前的位置
  RCLCPP_INFO(rclcpp::get_logger("rclcpp"), "Robot pose x : %.2fm; y : %.2fm; z : %.2fm", test_pose.position.x, test_pose.position.y, test_pose.position.z);

  // Shutdown ROS
  rclcpp::shutdown();
  spinner.join();  // <--- Join the thread before exiting
  return 0;
}
